
from files import FileTool
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
import numpy as np
from copy import deepcopy


def cal_angle(v1, v2):
    unit_vector_1 = v1 / np.linalg.norm(v1)
    unit_vector_2 = v2 / np.linalg.norm(v2)
    dot_product = np.dot(unit_vector_1, unit_vector_2)
    angle = np.degrees(np.arccos(dot_product))
    return angle


def calculate_control_pt(centerline, step=5, angle=150):
    """
    设置控制点，方便用户编辑中心线。
    目前是每隔固定的步长进行取点，判断每个点的曲率，如果大于设定阈值，则需添加控制点
    :centerline: 中心线坐标 （np.array）
    :step: 步长
    :angle: 夹角阈值
    :return: control_pts: 控制点的坐标(np.array),
             control_pts_index: 控制点在预处理后的中心线上的索引 (list)
    """
    control_pt_list = []
    control_pt_list.append(centerline[0])
    control_pts_index = [0]
    for i in range(1, (centerline.shape[0]) - 1, 1):
        v_ab = centerline[i - 1] - centerline[i]
        v_cb = centerline[i + 1] - centerline[i]
        angle_value = cal_angle(v_ab, v_cb)
        if abs(angle_value) < angle:
            control_pt_list.append(centerline[i])
            control_pts_index.append(i)
        if i % step == 0:
            control_pt_list.append(centerline[i - 1])
            control_pts_index.append(i - 1)
    if not np.max(np.isin(centerline[-1], control_pt_list)):
        control_pt_list.append(centerline[-1])
        control_pts_index.append(centerline.shape[0] - 1)
    control_pts = np.array(control_pt_list)
    control_pts_index = np.array(control_pts_index)
    # fig = plt.figure()
    # ax = fig.gca(projection='3d')
    # ax.plot(control_pts[:, 2], control_pts[:, 1], control_pts[:, 0], 'b*')
    # ax.plot(centerline[:, 2], centerline[:, 1], centerline[:, 0], 'r-')
    # plt.show()
    return control_pts, control_pts_index


class LabelFileGenerator:
    def __init__(self, filepath) -> None:
        """filepath 模型结果文件路径"""
        # self.files_list = None
        base_dir = Path(__file__).parent
        self.filepath = Path(filepath)
        self.template = FileTool.read_json(
            base_dir.joinpath('template.json').as_posix())

    def extra_cond(self, p):
        return p.stem == 'payload'

    def get_file(self, p, extra_cond=None):
        if p.is_file():
            if extra_cond and self.extra_cond(p) or not extra_cond:
                if p.suffix == '.json':
                    content = FileTool.read_json(p.resolve().as_posix())
                elif p.suffix == '.pkl':
                    content = FileTool.read_pkl(p.resolve().as_posix())
                else:
                    return
                file_info = [p.resolve(), content]
                yield file_info
        elif p.is_dir():
            for f in p.iterdir():
                generator = self.get_file(f, extra_cond)
                #递归时的生成器也需要通过yield返回
                for i in generator:
                    try:
                        yield i
                    except:
                        pass

    def append_cont(self, cont, name, res, step=20, angle=80):
        if not isinstance(cont, np.ndarray):
            cont = np.array(cont)
        choonsed_cont, _ = calculate_control_pt(cont, step=step, angle=angle)
        res['shapes'].append(
            {
                "label": name,
                "line_color": None,
                "fill_color": None,
                "points": choonsed_cont.tolist(),
                "shape_type": "polygon",
                "flags": {}
            })

    def append_point_line(self, point, name, type, res):
        res['shapes'].append(
            {
                "label": name,
                "line_color": None,
                "fill_color": None,
                "points": point,
                "shape_type": type,
                "flags": {}
            })

    def gene_4c_plax_annotations(self, payload, stem, annotation_dir):
        print(annotation_dir)
        for i, seg in payload.get('seg', {}).items():
            sub_val = seg.get("chambers", {})
            # chambers
            if isinstance(sub_val, dict):
                res = deepcopy(self.template)
                for k, v in sub_val.items():
                    # k=lv/la/ao/la
                    cont = v.get("cont", [])
                    if not cont:
                        print(i, 'is valid skip to generate annotations')
                        continue
                    if k == 'lv':
                        name = 'lvendo'
                        if v.get('lvpw_cont'):
                            self.append_cont(v['lvpw_cont'], 'lvpw', res,30)
                        if v.get('epi_cont'):
                            self.append_cont(v['epi_cont'], 'lvepi', res,30)
                        if v.get('ivs_cont'):
                            self.append_cont(v['ivs_cont'], 'ivs', res,30)
                    else:
                        name = k
                    self.append_cont(cont, name, res, 30)
                res['imagePath'] = stem+"_"+i+'.png'
                res['imageHeight'] = payload['img_size'][0]
                res['imageWidth'] = payload['img_size'][1]
                annotation_name = annotation_dir.as_posix()+("_"+i+'.json')
                FileTool.write_json(res, annotation_name)

    def gene_mm_annotations(self, payload, stem, annotation_dir):
        for i, seg in payload.get('seg', {}).items():
            res = deepcopy(self.template)
            res['imagePath'] = stem+"_"+i+'.png'
            res['imageHeight'] = payload['img_size'][0]
            res['imageWidth'] = payload['img_size'][1]
            regions = seg.get("regions", {})
            for pos, val in regions.items():
                if pos != 'lv':
                    if val.get('cont'):
                        cont = val['cont']
                        self.append_cont(
                            val['cont'], pos, res, step=20, angle=80)
                        cont_info = res['shapes'].pop()
                        points = np.array(cont_info['points'])
                        #找到极值点，拆分数组
                        min_v, max_v = points[:, 0].min(), points[:, 0].max()
                        left = np.argwhere(points[:, 0] == min_v)[0][0]
                        right = np.argwhere(
                            points[:, 0] == max_v)[0][0]
                        if left > right:
                            point_1 = points[right-1:left]
                            point_2 = np.vstack(
                                (points[left:], points[:right]))
                        else:
                            point_1 = points[left:right]
                            point_2 = np.vstack(
                                (points[right:],points[:left]))
                        # point_1 = np.sort(point_1,0)
                        # point_2 = np.sort(point_2,0)
                        #和更大者为2
                        if np.sum(point_1[:,1])>np.sum(point_2[:,1]):
                            point_1,point_2 = point_2,point_1
                        cont_info["shape_type"] = 'linestrip'
                        info = deepcopy(cont_info)
                        info['points'] = point_1.tolist()
                        info['label'] = pos+"1"
                        res['shapes'].append(info)
                        info = deepcopy(cont_info)
                        info['label'] = pos+'2'
                        info['points'] = point_2.tolist()
                        res['shapes'].append(info)
                # elif pos
            annotation_name = annotation_dir.as_posix()+("_"+i+'.json')
            FileTool.write_json(res, annotation_name)

    def gene_cw_pw_annotations(self, payload, stem, annotation_dir, view):
        # continue
        for i, seg in payload.get('seg', {}).items():
            res = deepcopy(self.template)
            res['imagePath'] = stem+"_"+i+'.png'
            res['imageHeight'] = payload['img_size'][0]
            res['imageWidth'] = payload['img_size'][1]
            for key in ["waves", "fwaves"]:
                # 1.2.840.113663.1500.1.423176588.3.22.20200731.161939.571_0 a4c-mv-cw
                vals = seg.get(key, [])
                for val in vals:
                    if val.get('cont'):
                        if view == 'A4C-TV-CW':
                            if val['pixel_Peak_point'][1] < val['anchor_rightdown_wh'][1]:
                                self.append_cont(
                                    val['cont'], 'up', res)
                            else:
                                self.append_cont(
                                    val['cont'], 'down', res)
                        else:
                            self.append_cont(val['cont'], 'PW', res)

            annotation_name = annotation_dir.as_posix()+("_"+i+'.json')
            FileTool.write_json(res, annotation_name)

    def main(self, outdir=None):
        """outdir 注释文件输出路径，不提供时默认文件生成路径为模型所在位置"""
        iter_files = self.get_file(self.filepath, extra_cond=True)
        for file in iter_files:
            fn = file[0]
            content = file[1]
            payload = content['ifs_payload']
            annotation_dir = Path(outdir) if outdir else fn.parent
            stem = Path(payload.get('file_path', '')).stem
            annotation_dir = annotation_dir.joinpath(stem)

            view = payload['view']
            # if view not in ['PLAX-2D', 'PLAX-LV-M', 'A4C-LW-TDI', 'A4C-IVS-TDI', 'A4C-TV-TDI']:
            #     continue
            print(stem, view)
            if view in ['A4C-2D', 'A2C-2D', 'PLAX-2D']:
                self.gene_4c_plax_annotations(payload, stem, annotation_dir)
            elif view in ['PLAX-LV-M']:
                self.gene_mm_annotations(payload, stem, annotation_dir)

            elif view in ['A4C-TV-PW', 'A4C-TV-CW']:
                self.gene_cw_pw_annotations(payload, stem, annotation_dir,view)
            elif view in ['A4C-LW-TDI', 'A4C-IVS-TDI', 'A4C-TV-TDI']:
                for i, seg in payload.get('seg', {}).items():
                    res = deepcopy(self.template)
                    res['imagePath'] = stem+"_"+i+'.png'
                    res['imageHeight'] = payload['img_size'][0]
                    res['imageWidth'] = payload['img_size'][1]
                    for key in ["waves", "fwaves"]:
                        # 1.2.840.113663.1500.1.423176588.3.22.20200731.161939.571_0 a4c-mv-cw
                        # 1.2.840.113663.1500.1.423176588.3.27.20200731.162032.678 A4C-LW-TDI
                        vals = seg.get(key, [])
                        for val in vals:
                            # points = []
                            for name in ['s2', 'e2', 'a2']:
                                if val.get(name):
                                    # points.append(val[name])
                                    self.append_point_line(
                                        [val[name]], name[0]+"'", 'point', res)
                        if seg.get('curve'):
                            cont, _ = calculate_control_pt(
                                np.array(seg['curve']), 20, 80)
                            self.append_point_line(
                                cont.tolist(), 'tdi', 'linestrip', res)
                    annotation_name = annotation_dir.as_posix()+("_"+i+'.json')
                    FileTool.write_json(res, annotation_name)


if __name__ == '__main__':
    plax_2d = '/Users/hongwei/workspace/label_repo/model_res/biomind_1/biomind/2478346d-cc3f-11ec-9bf4-0f15fdb8a48e'

    generator = LabelFileGenerator('/Users/hongwei/workspace/label_repo/model_res/'
                                   )
    generator.main('/Users/hongwei/workspace/label_repo/imgs')
